
from PyQt5.QtWidgets import ( QMainWindow, QWidget, QVBoxLayout,
                             QHBoxLayout, QLabel, QLineEdit, QPushButton,
                             QTextEdit, QTabWidget, QGroupBox, QProgressBar,
                             QSpinBox, QMessageBox, QRadioButton)
from PyQt5.QtCore import Qt, QThread, pyqtSignal
from brute_force import BruteForce  # 确保该文件存在
from S_DES import SDES  # 导入上面修改的SDES类


class BruteForceThread(QThread):
    """暴力破解线程"""
    finished = pyqtSignal(list, float)

    def __init__(self, plaintext, ciphertext, num_threads=4):
        super().__init__()
        self.plaintext = plaintext
        self.ciphertext = ciphertext
        self.num_threads = num_threads

    def run(self):
        brute_force = BruteForce(self.plaintext, self.ciphertext)
        keys = brute_force.crack_multi_thread(self.num_threads)
        elapsed = brute_force.get_elapsed_time()
        self.finished.emit(keys, elapsed)


class SDESSimpleGui(QMainWindow):
    """带模式选择的S-DES GUI"""

    def __init__(self):
        super().__init__()
        self.init_ui()

    def init_ui(self):
        """初始化界面"""
        self.setWindowTitle('S-DES 加密解密系统')
        self.setGeometry(100, 100, 800, 600)

        # 中心部件与主布局
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        main_layout = QVBoxLayout(central_widget)

        # 标签页
        tab_widget = QTabWidget()
        tab1 = self.create_encrypt_decrypt_tab()  # 加解密标签页
        tab2 = self.create_brute_force_tab()  # 暴力破解标签页
        tab_widget.addTab(tab1, '基础加解密')
        tab_widget.addTab(tab2, '暴力破解')

        main_layout.addWidget(tab_widget)

    def create_encrypt_decrypt_tab(self):
        """创建加解密标签页（包含模式选择）"""
        widget = QWidget()
        layout = QVBoxLayout(widget)

        # 1. 模式选择
        mode_group = QGroupBox('模式选择')
        mode_layout = QHBoxLayout()

        self.binary_mode_radio = QRadioButton("二进制模式")
        self.ascii_mode_radio = QRadioButton("ASCII模式")
        self.binary_mode_radio.setChecked(True)  # 默认二进制模式

        mode_layout.addWidget(self.binary_mode_radio)
        mode_layout.addWidget(self.ascii_mode_radio)
        mode_group.setLayout(mode_layout)
        layout.addWidget(mode_group)

        # 2. 密钥输入区
        key_group = QGroupBox('密钥设置（10位二进制）')
        key_layout = QHBoxLayout()
        self.key_input = QLineEdit()
        self.key_input.setPlaceholderText('示例：1010000010（仅允许0和1）')
        self.key_input.setMaxLength(10)
        key_layout.addWidget(QLabel('10-bit 密钥：'))
        key_layout.addWidget(self.key_input, stretch=1)
        key_group.setLayout(key_layout)
        layout.addWidget(key_group)

        # 3. 加密区
        encrypt_group = QGroupBox('加密功能')
        encrypt_layout = QVBoxLayout()

        # 明文输入（根据模式显示不同提示）
        plain_layout = QHBoxLayout()
        self.plaintext_input = QTextEdit()
        self.plaintext_input.setPlaceholderText('示例：10101010（8位二进制）')
        self.plaintext_input.setMaximumHeight(60)
        # 切换模式时改变输入提示
        self.binary_mode_radio.toggled.connect(lambda: self.update_placeholders())
        plain_layout.addWidget(QLabel('输入明文：'))
        plain_layout.addWidget(self.plaintext_input, stretch=1)
        encrypt_layout.addLayout(plain_layout)

        # 加密按钮
        encrypt_btn = QPushButton('执行加密')
        encrypt_btn.clicked.connect(self.encrypt)
        encrypt_layout.addWidget(encrypt_btn, alignment=Qt.AlignRight)

        # 密文输出
        cipher_output_layout = QHBoxLayout()
        self.ciphertext_output = QTextEdit()
        self.ciphertext_output.setReadOnly(True)
        self.ciphertext_output.setStyleSheet('background-color: #f0f0f0;')
        self.ciphertext_output.setMaximumHeight(60)
        cipher_output_layout.addWidget(QLabel('加密结果：'))
        cipher_output_layout.addWidget(self.ciphertext_output, stretch=1)
        encrypt_layout.addLayout(cipher_output_layout)

        encrypt_group.setLayout(encrypt_layout)
        layout.addWidget(encrypt_group)

        # 4. 解密区
        decrypt_group = QGroupBox('解密功能')
        decrypt_layout = QVBoxLayout()

        # 密文输入
        cipher_input_layout = QHBoxLayout()
        self.ciphertext_input = QTextEdit()
        self.ciphertext_input.setPlaceholderText('示例：11001010（8位二进制）')
        self.ciphertext_input.setMaximumHeight(60)
        cipher_input_layout.addWidget(QLabel('输入密文：'))
        cipher_input_layout.addWidget(self.ciphertext_input, stretch=1)
        decrypt_layout.addLayout(cipher_input_layout)

        # 解密按钮
        decrypt_btn = QPushButton('执行解密')
        decrypt_btn.clicked.connect(self.decrypt)
        decrypt_layout.addWidget(decrypt_btn, alignment=Qt.AlignRight)

        # 明文输出
        plain_output_layout = QHBoxLayout()
        self.plaintext_output = QTextEdit()
        self.plaintext_output.setReadOnly(True)
        self.plaintext_output.setStyleSheet('background-color: #f0f0f0;')
        self.plaintext_output.setMaximumHeight(60)
        plain_output_layout.addWidget(QLabel('解密结果：'))
        plain_output_layout.addWidget(self.plaintext_output, stretch=1)
        decrypt_layout.addLayout(plain_output_layout)

        decrypt_group.setLayout(decrypt_layout)
        layout.addWidget(decrypt_group)

        layout.addStretch(1)
        return widget

    def create_brute_force_tab(self):
        """创建暴力破解标签页"""
        # 此处代码与之前类似，保持不变
        widget = QWidget()
        layout = QVBoxLayout(widget)

        # 1. 明密文对输入区
        input_group = QGroupBox('已知明密文对（均为8位二进制）')
        input_layout = QVBoxLayout(input_group)

        # 明文输入
        plain_layout = QHBoxLayout()
        self.bf_plaintext_input = QLineEdit()
        self.bf_plaintext_input.setPlaceholderText('示例：10101010（明文）')
        self.bf_plaintext_input.setMaxLength(8)
        plain_layout.addWidget(QLabel('8-bit 明文：'))
        plain_layout.addWidget(self.bf_plaintext_input, stretch=1)
        input_layout.addLayout(plain_layout)

        # 密文输入
        cipher_layout = QHBoxLayout()
        self.bf_ciphertext_input = QLineEdit()
        self.bf_ciphertext_input.setPlaceholderText('示例：11001010（密文）')
        self.bf_ciphertext_input.setMaxLength(8)
        cipher_layout.addWidget(QLabel('8-bit 密文：'))
        cipher_layout.addWidget(self.bf_ciphertext_input, stretch=1)
        input_layout.addLayout(cipher_layout)

        # 线程数设置
        threads_layout = QHBoxLayout()
        self.thread_spinbox = QSpinBox()
        self.thread_spinbox.setRange(1, 16)
        self.thread_spinbox.setValue(4)
        threads_layout.addWidget(QLabel('破解线程数：'))
        threads_layout.addWidget(self.thread_spinbox)
        threads_layout.addStretch(1)
        input_layout.addLayout(threads_layout)

        input_group.setLayout(input_layout)
        layout.addWidget(input_group)

        # 2. 破解控制区
        control_layout = QHBoxLayout()
        self.bf_start_btn = QPushButton('开始暴力破解')
        self.bf_start_btn.clicked.connect(self.start_brute_force)
        control_layout.addWidget(self.bf_start_btn)

        self.bf_progress = QProgressBar()
        self.bf_progress.setRange(0, 0)
        self.bf_progress.hide()
        control_layout.addWidget(self.bf_progress, stretch=1)
        layout.addLayout(control_layout)

        # 3. 破解结果显示区
        result_group = QGroupBox('破解结果')
        result_layout = QVBoxLayout(result_group)
        self.bf_result_output = QTextEdit()
        self.bf_result_output.setReadOnly(True)
        result_layout.addWidget(self.bf_result_output)
        result_group.setLayout(result_layout)
        layout.addWidget(result_group, stretch=1)

        return widget

    # ---------------------- 辅助方法 ----------------------
    def update_placeholders(self):
        """根据选择的模式更新输入框提示文字"""
        if self.binary_mode_radio.isChecked():
            self.plaintext_input.setPlaceholderText('示例：10101010（8位二进制，可多个用空格分隔）')
            self.ciphertext_input.setPlaceholderText('示例：11001010（8位二进制，可多个用空格分隔）')
        else:
            self.plaintext_input.setPlaceholderText('示例：Hello World（ASCII字符串）')
            self.ciphertext_input.setPlaceholderText('示例：加密后的ASCII字符串')

    def validate_key(self, key_str):
        """验证密钥：10位二进制"""
        if len(key_str) != 10:
            QMessageBox.warning(self, '输入错误', '密钥必须是10位二进制数！')
            return False
        if not all(c in '01' for c in key_str):
            QMessageBox.warning(self, '输入错误', '密钥只能包含数字0和1！')
            return False
        return True

    def validate_binary_input(self, input_str, desc):
        """验证二进制输入"""
        blocks = input_str.split()
        for block in blocks:
            if len(block) != 8:
                QMessageBox.warning(self, '输入错误', f'{desc}必须是8位二进制数！')
                return False
            if not all(c in '01' for c in block):
                QMessageBox.warning(self, '输入错误', f'{desc}只能包含数字0和1！')
                return False
        return True

    # ---------------------- 核心功能逻辑 ----------------------
    def encrypt(self):
        """执行加密（根据模式选择相应方法）"""
        key = self.key_input.text().strip()
        if not self.validate_key(key):
            return

        plaintext = self.plaintext_input.toPlainText().strip()
        if not plaintext:
            QMessageBox.warning(self, '输入错误', '请输入明文！')
            return

        try:
            sdes = SDES(key)
            result = ""

            if self.binary_mode_radio.isChecked():
                # 二进制模式加密
                if not self.validate_binary_input(plaintext, '明文'):
                    return

                # 处理多个块
                blocks = plaintext.split()
                encrypted_blocks = []
                for block in blocks:
                    encrypted = sdes.encrypt_block(block)
                    encrypted_blocks.append(encrypted)
                result = ' '.join(encrypted_blocks)

            else:
                # ASCII模式加密
                result = sdes.encrypt_ascii(plaintext)

            self.ciphertext_output.setText(result)
            QMessageBox.information(self, '加密成功', '加密完成！')

        except Exception as e:
            QMessageBox.critical(self, '加密失败', f'错误原因：{str(e)}')

    def decrypt(self):
        """执行解密（根据模式选择相应方法）"""
        key = self.key_input.text().strip()
        if not self.validate_key(key):
            return

        ciphertext = self.ciphertext_input.toPlainText().strip()
        if not ciphertext:
            QMessageBox.warning(self, '输入错误', '请输入密文！')
            return

        try:
            sdes = SDES(key)
            result = ""

            if self.binary_mode_radio.isChecked():
                # 二进制模式解密
                if not self.validate_binary_input(ciphertext, '密文'):
                    return

                # 处理多个块
                blocks = ciphertext.split()
                decrypted_blocks = []
                for block in blocks:
                    decrypted = sdes.decrypt_block(block)
                    decrypted_blocks.append(decrypted)
                result = ' '.join(decrypted_blocks)

            else:
                # ASCII模式解密
                result = sdes.decrypt_ascii(ciphertext)

            self.plaintext_output.setText(result)
            QMessageBox.information(self, '解密成功', '解密完成！')

        except Exception as e:
            QMessageBox.critical(self, '解密失败', f'错误原因：{str(e)}')

    def start_brute_force(self):
        """启动暴力破解"""
        plaintext = self.bf_plaintext_input.text().strip()
        ciphertext = self.bf_ciphertext_input.text().strip()
        num_threads = self.thread_spinbox.value()

        # 验证输入
        if not self.validate_binary_input(plaintext, '明文'):
            return
        if not self.validate_binary_input(ciphertext, '密文'):
            return

        # 初始化破解状态
        self.bf_start_btn.setEnabled(False)
        self.bf_progress.show()
        self.bf_result_output.clear()

        # 启动破解线程
        self.brute_thread = BruteForceThread(plaintext, ciphertext, num_threads)
        self.brute_thread.finished.connect(self.on_brute_force_finished)
        self.brute_thread.start()

    def on_brute_force_finished(self, keys, elapsed):
        """暴力破解完成回调"""
        self.bf_start_btn.setEnabled(True)
        self.bf_progress.hide()

        if not keys:
            result = f'破解失败：未找到匹配的10位密钥\n耗时：{elapsed:.2f}秒'
        else:
            result = f'破解成功！\n找到 {len(keys)} 个匹配密钥：\n'
            for i, key in enumerate(keys, 1):
                result += f'  {i}. 密钥：{key}\n'
            result += f'\n总耗时：{elapsed:.2f}秒'

        self.bf_result_output.setText(result)
        QMessageBox.information(self, '破解完成', result.split('\n')[0])



